import numpy as np
import sympy as sp
import sys

sys.path.append('../pde')
sys.path.append('../fem')



from fealpy.mesh import MeshFactory
from fealpy.mesh import TriangleMesh
from fealpy.pde.linear_elasticity_model import QiModel3d, PolyModel3d, Model2d, HuangModel2d
from linear_elasticity_model2D import GenLinearElasticitymodel2D
from LinearElasticityFEMModel import LinearElasticityFEMModel 
from fealpy.tools.show import showmultirate


import matplotlib.pyplot as plt


m = int(sys.argv[1])
p = int(sys.argv[2])
n = int(sys.argv[3])
s = int(sys.argv[4])

E = 1e3
nu = 0.3
lam = E*nu/((1+nu)*(1-2*nu))
mu = E/(2*(1+nu))


pi = sp.pi
sin = sp.sin
cos = sp.cos
exp = sp.exp
ln = sp.ln
mf = MeshFactory()


def plotshow(X,Y,Z,colorbar=None,s=0):
    '''
    绘图用
    '''
    from matplotlib.colors import Normalize
    if colorbar is None:
        colorbar = 'rainbow'

    extent = (0,1,0,1)
    
    x = np.min(Z)
    if np.abs(x)>10:
        vmin = int(x)
    else:
        vmin = x
    
    x = np.max(Z)
    if np.abs(x)>10:
        vmax = int(x)+1
    else:
        vmax = x
    print(np.min(Z),np.max(Z))
    
    #vmin = np.min(Z)
    #vmax = np.max(Z)

    norm = Normalize(vmin = vmin,vmax = vmax,clip=True)
    n = Z.shape[-1] #默认最后一个轴是要画的u的个数

    fig, axes = plt.subplots(ncols=n,nrows=1)
    if n == 1:
        if s == 0:
            im = axes.contourf(X,Y,Z[...,0],extent = extent,norm = norm,cmap='rainbow')
            for i in range(X.shape[0]):
                axes.plot(X[i],Y[i],linewidth=0.5,c='k')
            for j in range(Y.shape[1]):
                axes.plot(X[:,j],Y[:,j],linewidth=0.5,c='k')
        else:
            im = axes.imshow(Z[...,0],extent = extent,norm = norm,cmap='rainbow')
    elif n>1:
        for i in range(n):
            if i < n-1:
                if s == 0:
                    axes[i].contourf(X,Y,Z[...,i],extent = extent,norm = norm,cmap='rainbow')
                else:
                    axes[i].imshow(Z[...,i],extent = extent,norm = norm,cmap='rainbow')
            else:
                if s == 0:
                    im = axes[i].contourf(X,Y,Z[...,i],extent = extent,norm = norm,cmap='rainbow')
                else:
                    im = axes[i].imshow(Z[...,i],extent = extent,norm = norm,cmap='rainbow')
            for k in range(X.shape[0]):
                axes[i].plot(X[k],Y[k],linewidth=0.5,c='k')
            for l in range(Y.shape[1]):
                axes[i].plot(X[:,l],Y[:,l],linewidth=0.5,c='k')

    fig.subplots_adjust(right=0.8)
    l = 0.82
    b = 0.26
    w = 0.015
    h = 0.47
    rect = [l,b,w,h]
    cbar_ax = fig.add_axes(rect) 
    fig.colorbar(im,cax=cbar_ax)







# 给定位移算例
#pde = TestDirichlemodel(lam=lam,mu=mu,stress_direction='stress_nt',displacement_direction='displacement_n')

#给定应力算例

#一般pde算例
x = sp.symbols('x0:2')
#u = [0.01*(1-x[0]),0]
#u = [sin(2*pi*x[0])*sin(2*pi*x[1]),sin(2*pi*x[0])*sin(2*pi*x[1])]
u = [exp(x[0]+x[1]),exp(x[0]-x[1])]
#u = [0,x[1]]
#u = [0,(x[0]-1)**7*(x[1]-1)**7]
#u = [-(1-x[0])*ln(1.5-x[0]),-(1-x[0])*ln(1.5-x[1])]
#u = [sin(2*pi*x[0])*sin(2*pi*x[1]),sin(2*pi*x[0])*sin(2*pi*x[1])]
print(len(u))
lam = 1
mu = 0.5

idx = 1
if idx == 1:

    pde = GenLinearElasticitymodel2D(u,x,lam=lam,mu=mu,
            Dirichletbd_n='(x0==1)|(x0==0)|(x1==0)|(x1==1)',
            Dirichletbd_t='(x0==1)|(x0==0)|(x1==0)|(x1==1)')
elif idx == 1:

    pde = GenLinearElasticitymodel2D(u,x,lam=lam,mu=mu,
            Dirichletbd_n='(x0==1)|(x1==0)',Dirichletbd_t='(x0==1)|(x1==0)',
            Neumannbd_nn='(x0==0)|(x1==1)',Neumannbd_nt='(x0==0)|(x1==1)')

else:
     pde = GenLinearElasticitymodel2D(u,x,lam=lam,mu=mu,
            Dirichletbd_n='(x0==1)|(x1==0)|(x0==0)|(x1==1)',
            Dirichletbd_t='(x0==1)|(x1==0)|(x0==0)|(x1==1)')


#pde = DisplacementTestmodel(lam=lam,mu=mu)
#pde = StressTestmodel(lam=lam,mu=mu)
#pde = Stress_concentrationTestmodel(lam=lam,mu=mu)



node = np.array([[0,0],[0,1],[1,0],[1,1],[1/2,1/2]],dtype=np.float)
cell = np.array([[0,2,4],[3,4,2],[1,4,3],[0,4,1]],dtype=np.int)
mesh = TriangleMesh(node=node,cell=cell)
mesh.uniform_refine(n)
#print(np.sqrt(mesh.number_of_cells()))
integrator = mesh.integrator(5)



#mesh = mf.boxmesh2d(pde.domain(),nx=n,ny=n,meshtype='tri')

showmesh = False
if showmesh:
    fig = plt.figure()
    axes = fig.gca()
    mesh.add_plot(axes)
    mesh.find_node(axes, showindex=True)
    mesh.find_cell(axes, showindex=True)
    plt.show()

if m == 1:
   
    
    maxit = 5
    errorType = ['$||\sigma - \sigma_h ||_{0}$',
                '$||div(\sigma - \sigma_h)||_{0}$',
                '$||u - u_h||_{0}$',
                '$||\sigma - \sigma_I ||_{0}$',
                '$||div(\sigma - \sigma_I)||_{0}$',
                '$||u - u_I||_{0}$',
                '$|| u_I - u_h||_{0}$'
                ]
    Ndof = np.zeros((maxit,))
    Internumber = np.zeros((maxit,))
    solve_time = np.zeros((maxit,))
    linear_operator_time = np.zeros((maxit,))
    errorMatrix = np.zeros((len(errorType), maxit), dtype=np.float)


    for i in range(maxit):
        fem = LinearElasticityFEMModel(mesh, pde, p, integrator)
        fem.get_AA_and_bb()
        #fem.solve()
        #fem.fast_solve()
        fem.matlab_solve()
        print('\n')
        fem.Finite_emelent_solution()
        tgdof = fem.tensorspace.number_of_global_dofs()
        vgdof = fem.vectorspace.number_of_global_dofs()
        gdof = tgdof + vgdof
        Ndof[i] = gdof
        #Internumber[i] = fem.internumber
        #solve_time[i] = fem.solve_time
        #linear_operator_time[i] = fem.linear_operator_time
        errorMatrix[:, i] = fem.error()
        if i < maxit - 1:
            mesh.uniform_refine()
            
    print('Ndof:', Ndof)
    #print('Inter number:', Internumber)
    #print('solve time:',solve_time)
    #print('linear operator time:',linear_operator_time)
    print('error:', errorMatrix)
    showmultirate(plt, 0, Ndof, errorMatrix, errorType)
    plt.show()



elif m == 2:
    fem = LinearElasticityFEMModel(mesh, pde, p, integrator)
    fem.get_AA_and_bb()
    fem.Save_A_and_b_to_matlab()



























elif m==3: #绘图
    
    show_displacement = True
    show_stress = False
    from Cartesian_coordinates_function import Cartesian_coordinates_function

    fem = LinearElasticityFEMModel(mesh, pde, p, integrator)
    fem.solve()
    #fem.fast_solve()
    sh = fem.sh
    uh = fem.uh
    sh_val = sh.value
    uh_val = uh.value
    tensorspace = fem.tensorspace
    vectorspace = fem.vectorspace
    sh_car_func = Cartesian_coordinates_function(tensorspace,sh[:])
    uh_car_func = Cartesian_coordinates_function(vectorspace,uh[:])

    
    
    x = np.linspace(0,1,50)
    y = np.linspace(1,0,50) #用imshows绘图时,y坐标轴刚好反了过来
    X,Y = np.meshgrid(x,y)
    shape = X.shape
    shape +=(2,)
    p = np.zeros(shape,dtype=float)
    p[...,0] = X
    p[...,1] = Y
    
    
    
    if show_displacement:
        shape = X.shape
        # 绘制位移方向图

        shape +=(1,)
        Z = np.zeros(shape,dtype=np.float) 
        displacementh = uh_car_func.value(p)  
        for i in range(2):
            if 'displacement' in dir(pde):
                Z[...,0] = displacementh[...,i]
                #Z[...,0] = pde.displacement(p)[...,i]
                #Z[...,0] = np.abs(displacementh[...,i]-pde.displacement(p)[...,i])/np.max(np.abs(pde.displacement(p)[...,i]))
            else:
                Z[...,0] = displacementh[...,i]
            plotshow(X,Y,Z,s=s)
        
        #plt.show()




    
    if show_stress: 
        # 绘制应力图
        shape = X.shape
        sgimah = sh_car_func.values(p)
        Recover_sgimah = uh_car_func.Recover_sgimah(p,lam=lam,mu=mu)

 
        shape +=(1,)
        Z = np.zeros(shape,dtype=np.float)    
        for i in range(3):
            if 'stress' in dir(pde):
                Z[...,0] = sgimah[...,i]
                #Z[...,0] = pde.stress(p)[...,[0,1,1],[0,1,0]][...,i]
                #Z[...,0] = np.abs(sgimah[...,i]-(pde.stress(p)[...,[0,1,1],[0,1,0]])[...,i])/np.max((np.abs((pde.stress(p)[...,[0,1,1],[0,1,0]])[...,i])))
            else:
                Z[...,0] = sgimah[...,i]
            plotshow(X,Y,Z,s=s)

        plt.show()


elif m == 4:
    from Cartesian_coordinates_function import Cartesian_coordinates_function
    fem = LinearElasticityFEMModel(mesh, pde, p, integrator)
    fem.solve()
    sh = fem.sh
    uh = fem.uh
    sh_val = sh.value
    uh_val = uh.value

    tensorspace = fem.tensorspace
    vectorspace = fem.vectorspace
    
    sh_car_func = Cartesian_coordinates_function(tensorspace,sh[:])
    uh_car_func = Cartesian_coordinates_function(vectorspace,uh[:])

    y = np.linspace(0,1,4)
    x = np.zeros(y.shape)
    p = np.zeros(y.shape+(2,))
    p[:,0] = x
    p[:,1] = y
    #print(sh_car_func.values(p))
    #print(pde.stress(p)[...,[0,1,1],[0,1,0]])



   